"""基于论文[1] Park Woosang, Jung Youeyun, Bang Hyochoong, et al. Robust crater triangle matching algorithm for planetary landing navigation[J]. Journal of Guidance, Control, and Dynamics, 2019, 42(2): 402-410

"""

from .base import TriadPlainBuilder
from utils.ellipse import ellipse_args, radius_ellipse
import numpy as np
from itertools import combinations


class ParkBuilder(TriadPlainBuilder):
    def invariant(self, i, j, k, Q1, Q2, Q3) -> tuple[np.ndarray, np.ndarray]:
        N = Q3.shape[0]
        Q1 = Q1[None].repeat(N, axis=0)
        Q2 = Q2[None].repeat(N, axis=0)
        d1 = np.mean(radius_ellipse(Q1), axis=0)
        d2 = np.mean(radius_ellipse(Q2), axis=0)
        d3 = np.mean(radius_ellipse(Q3), axis=0)
        # 按直径从小到大排序
        ind = np.argsort((d1, d2, d3), axis=0)
        Q = np.array((Q1, Q2, Q3))[ind, np.arange(N)]
        x1, y1, a1, b1, _ = ellipse_args(Q[0])
        x2, y2, a2, b2, _ = ellipse_args(Q[1])
        x3, y3, a3, b3, _ = ellipse_args(Q[2])
        d1 = a1 + b1
        d2 = a2 + b2
        d3 = a3 + b3
        # 三角形的三个顶点坐标
        A = np.array((x1, y1, np.ones_like(x1))).T
        B = np.array((x2, y2, np.ones_like(x1))).T
        Cs = np.array((x3, y3, np.ones_like(x1))).T
        AC = np.cross(A, Cs, axis=1)
        BC = np.cross(B, Cs, axis=1)
        AB = np.cross(A, B, axis=1)
        # 求解交点
        Q = np.vstack((Q[0], Q[1], Q[2], Q[1], Q[2], Q[0]))
        M = np.vstack((AB, BC, AC, AB, BC, AC))
        # 由于M、Q由六个点组成，因此当一个交点为nan时应当全部舍去
        ## view(-1, 6, 4)分别代表，6个交点，num_tri代表输入的被拉平的三角形个数，但是为了保证
        ## 每个三角形都不存在Nan，因此将拉平结果还原为原始维度，然后判断每个三角形内是否有Nan
        cross = self.solve_cross_points(M, Q)
        N = cross.shape[0] // 2
        cross = self.judge_inside(
            np.stack((A, B, Cs), axis=0),
            np.stack((B, Cs, A), axis=0),
            cross[:N].reshape(3, -1, 2, 3),
            cross[N:].reshape(3, -1, 2, 3),
        )
        valid_ind = ~np.isnan(cross).any(axis=(0, 2, 3))
        # 计算不变量
        I = self.g_func(cross.transpose(1, 0, 2, 3).reshape(N // 3, -1, 3)).T
        ijk = np.vstack((np.ones_like(k) * i, np.ones_like(k) * j, k))
        # 按直径从小到大排序
        ijk = np.take_along_axis(ijk, ind, axis=0)
        return ijk[:, valid_ind], I[:, valid_ind]

    def solve_cross_points(self, m, Q):
        """求解直线m与椭圆Q的交点，以齐次形式返回"""
        assert Q.shape[0] == m.shape[0]
        ind = np.abs(m[:, 0]) > 1e6 * np.abs(m[:, 1])
        res = ~ind
        cross = np.zeros((3, 2, m.shape[0]))
        ## # 对于一般情况，用通用方法求解交点
        t_n = m[res, 2] / m[res, 1]
        m_n = m[res, 0] / m[res, 1]
        a = Q[res, 0, 0] - 2 * Q[res, 0, 1] * m_n + Q[res, 1, 1] * m_n**2
        b = 2 * (
            -Q[res, 0, 1] * t_n
            + Q[res, 1, 1] * m_n * t_n
            + Q[res, 0, 2]
            - Q[res, 1, 2] * m_n
        )
        c = Q[res, 1, 1] * t_n**2 - 2 * Q[res, 1, 2] * t_n + Q[res, 2, 2]
        ## 待解矩阵
        x = np.array(
            (
                (-b + np.sqrt(b**2 - 4 * a * c)) / 2 / a,
                (-b - np.sqrt(b**2 - 4 * a * c)) / 2 / a,
            )
        )
        y = -m_n * x - t_n
        cross[..., res] = np.array((x, y, np.ones_like(x)))
        # if (np.isclose(x[0], x[1]) & np.isclose(y[0], y[1])).sum() > 0:
        #     print("error")
        ## # 对于特殊情况，直线斜率为0，直接求解
        if ind.sum() > 0:
            x = -m[ind, 2] / m[ind, 0]
            delta = np.sqrt(
                (Q[ind, 1, 2] + Q[ind, 0, 1] * x) ** 2
                - Q[ind, 1, 1]
                * (Q[ind, 2, 2] + 2 * Q[ind, 0, 2] * x + Q[ind, 0, 0] * x**2)
            )
            b = Q[ind, 1, 2] + Q[ind, 0, 1] * x
            x = np.array((x, x))
            y = np.array(
                (
                    (-b + delta) / Q[ind, 1, 1],
                    (-b - delta) / Q[ind, 1, 1],
                )
            )
            cross[..., ind] = np.array((x, y, np.ones_like(x)))
        return cross.T

    def judge_inside(self, B, C, B_cross, C_cross):
        # 判断在内侧
        B = B[..., None, :]
        C = C[..., None, :]
        BB_cross = B_cross - B
        CB_cross = B_cross - C
        CC_cross = C_cross - C
        BC_cross = C_cross - B
        ## 判断是否在内侧，用小阈值排除等于零的情况
        index_B = (CB_cross * BB_cross).sum(axis=3) < -1e-3
        index_C = (CC_cross * BC_cross).sum(axis=3) < -1e-3
        # ind_B和ind_C应当是相等的，因为它们的原始顺序是一致的
        ind = (index_B.sum(axis=2) == 1).all(axis=0) & (index_C.sum(axis=2) == 1).all(
            axis=0
        )
        index_B = index_B & ind[None, :, None]
        index_C = index_C & ind[None, :, None]
        ## 返回所得交点和公共交集
        result = np.ones_like(B_cross) * np.nan
        result[:, ind] = np.stack((B_cross[index_B], C_cross[index_C]), axis=1).reshape(
            3, -1, 2, 3
        )
        return result

    def g_func(self, x):
        assert x.shape[1] == 6
        # 获取五个点的组合
        comb = np.array(list(combinations(range(6), 5)))
        mu, nu = self.mu_nu(x[:, comb])
        # 计算不变量
        return np.concatenate((mu, nu), axis=1)

    def J(self, x):
        return (
            2 * x**6 - 6 * x**5 + 9 * x**4 - 8 * x**3 + 9 * x**2 - 6 * x + 2
        ) / (x**6 - 3 * x**5 + 3 * x**4 - x**3 + 3 * x**2 - 3 * x + 1)

    def mu_nu(self, x):
        assert x.shape[2] == 5
        mu = (
            np.linalg.det(x[:, :, [3, 2, 0]])
            * np.linalg.det(x[:, :, [4, 1, 0]])
            / np.linalg.det(x[:, :, [3, 1, 0]])
            / np.linalg.det(x[:, :, [4, 2, 0]])
        )
        nu = (
            np.linalg.det(x[:, :, [3, 1, 0]])
            * np.linalg.det(x[:, :, [4, 2, 1]])
            / np.linalg.det(x[:, :, [3, 2, 1]])
            / np.linalg.det(x[:, :, [4, 1, 0]])
        )
        return mu, nu
